import mxnet as mx
from mxnet.gluon import nn
import gluoncv
import mxnet.ndarray as nd
import numpy as np
import gluonnlp
from models.transformer import get_transformer_encoder_decoder
from mxnet import autograd as ag

class Conv(nn.HybridBlock):
    def __init__(self,norm = False,channels = 64,ks = 3,activation = "relu"):
        super(Conv,self).__init__()
        self.conv = nn.Conv2D(channels=channels,kernel_size=(ks,ks),padding=((ks-1)//2,(ks-1)//2),activation=activation)
        self.conv.initialize()
        self._norm = norm
        if self._norm:
            self.bn = nn.BatchNorm()
            self.bn.initialize()
    def hybrid_forward(self, F, x, *args, **kwargs):
        x = self.conv(x)
        if self._norm:
            x = self.bn(x)
        return x

class Lenet(mx.gluon.Block):
    def __init__(self):
        super(Lenet,self).__init__()
        self.feature = nn.Sequential()
        with self.feature.name_scope():
            self.feature.add(
                Conv(True,64),
                Conv(True, 128),
                Conv(True, 128),
                Conv(True, 256),
                Conv(True, 256),
                Conv(True, 512),
                Conv(True, 512),
            )
        print(self.feature)
    def forward(self, x):
        return self.feature(x)
class Resnet_CTC(mx.gluon.Block):
    def __init__(self,alphabet_size = 111,sequence_len = 1000):
        super(Resnet_CTC,self).__init__()
        self.features = Lenet()
        # self.ctc = mx.gluon.loss.CTCLoss(layout='NTC', label_layout='NT')
        self.conv = nn.Conv2D(channels=sequence_len , kernel_size=(1, 1), strides=(1, 1))
        self.conv.bias.initialize(init=mx.init.Zero())
        self.conv.weight.initialize(init=mx.init.Normal())

        self.dense = nn.Dense(512,flatten=False)
        self.dense.initialize()

        self._alphabet_size = alphabet_size
        self._sequence_len = sequence_len
        import mxnet.gluon as gluon

        self.encoder = mx.gluon.rnn.LSTM(hidden_size=self._alphabet_size,num_layers=1,layout='NTC')
        self.encoder.initialize()
        self.decoder = mx.gluon.rnn.GRU(hidden_size=self._alphabet_size,num_layers=1,layout='NTC')
        self.decoder.initialize()

    def forward_test(self, x):
        x = self.features(x)
        return x
    def forward(self, *args):

        return self.forward_test(*args)
if __name__ == '__main__':
    res = Resnet_CTC()
    res.collect_params().reset_ctx(mx.gpu(2))
    x= mx.nd.zeros(shape=(1,3,64,256),ctx=mx.gpu(2))
    y = res(x)
    print(y[0][0])
    print(y.shape)
